"""
C++ inheritance analyzer.

Provides comprehensive analysis of C++ class hierarchies including
inheritance relationships, virtual function resolution, and polymorphism.
"""

from typing import Dict, List, Any, Optional, Set, Tuple
from dataclasses import dataclass, field
from collections import defaultdict, deque
import logging

from ...core.error_handler import log_debug, log_info


@dataclass
class InheritanceRelation:
    """Represents an inheritance relationship."""
    derived_class: str
    base_class: str
    access_level: str  # public, protected, private
    is_virtual: bool = False
    location: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ClassInfo:
    """Enhanced class information for inheritance analysis."""
    name: str
    qualified_name: str
    location: Dict[str, Any]
    is_abstract: bool = False
    is_final: bool = False
    base_classes: List[InheritanceRelation] = field(default_factory=list)
    derived_classes: List[str] = field(default_factory=list)
    virtual_methods: List[Dict[str, Any]] = field(default_factory=list)
    pure_virtual_methods: List[Dict[str, Any]] = field(default_factory=list)
    overridden_methods: List[Dict[str, Any]] = field(default_factory=list)
    polymorphic_calls: List[Dict[str, Any]] = field(default_factory=list)
    vtable_size: int = 0
    inheritance_depth: int = 0


@dataclass
class VirtualMethodInfo:
    """Information about virtual methods."""
    name: str
    qualified_name: str
    declaring_class: str
    overriding_classes: List[str] = field(default_factory=list)
    is_pure_virtual: bool = False
    is_final: bool = False
    override_chain: List[str] = field(default_factory=list)
    location: Dict[str, Any] = field(default_factory=dict)


class InheritanceAnalyzer:
    """
    Analyzes C++ inheritance hierarchies and virtual function behavior.
    
    Provides comprehensive analysis of class hierarchies, virtual function
    resolution, polymorphic behavior, and vtable construction.
    """
    
    def __init__(self):
        """Initialize the inheritance analyzer."""
        self.reset()
        log_debug("Initialized C++ inheritance analyzer")
    
    def reset(self):
        """Reset analyzer state for new analysis."""
        self.class_hierarchy = {}
        self.inheritance_relations = []
        self.virtual_methods = {}
        self.abstract_classes = set()
        self.final_classes = set()
        self.diamond_inheritance_cases = []
        self.vtables = {}
        
        # Analysis caches
        self.class_index = {}
        self.method_index = defaultdict(list)
        self.inheritance_graph = defaultdict(set)
        self.reverse_inheritance_graph = defaultdict(set)
        
        # Statistics
        self.stats = {
            'total_classes': 0,
            'abstract_classes': 0,
            'final_classes': 0,
            'inheritance_relations': 0,
            'virtual_methods': 0,
            'pure_virtual_methods': 0,
            'diamond_inheritance_cases': 0,
            'polymorphic_classes': 0,
            'max_inheritance_depth': 0,
            'avg_inheritance_depth': 0
        }
    
    def analyze_inheritance(self, ast_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Analyze inheritance relationships in C++ code.
        
        Args:
            ast_data: Complete AST data from parser
            
        Returns:
            Comprehensive inheritance analysis results
        """
        self.reset()
        
        log_info("Starting C++ inheritance analysis")
        
        # Build class index
        self._build_class_index(ast_data)
        
        # Extract inheritance relationships
        self._extract_inheritance_relations(ast_data)
        
        # Build inheritance graphs
        self._build_inheritance_graphs()
        
        # Analyze virtual methods
        self._analyze_virtual_methods(ast_data)
        
        # Detect diamond inheritance
        self._detect_diamond_inheritance()
        
        # Build virtual function tables
        self._build_vtables()
        
        # Calculate inheritance depths
        self._calculate_inheritance_depths()
        
        # Analyze polymorphic behavior
        self._analyze_polymorphic_behavior(ast_data)
        
        return self._format_inheritance_results()
    
    def _build_class_index(self, ast_data: Dict[str, Any]):
        """Build an index of all classes."""
        for cls in ast_data.get('classes', []):
            class_info = ClassInfo(
                name=cls['name'],
                qualified_name=cls['qualified_name'],
                location=cls['location'],
                is_abstract=cls.get('is_abstract', False),
                is_final=cls.get('is_final', False)
            )
            
            self.class_index[cls['qualified_name']] = class_info
            self.class_hierarchy[cls['qualified_name']] = class_info
            
            if class_info.is_abstract:
                self.abstract_classes.add(cls['qualified_name'])
            if class_info.is_final:
                self.final_classes.add(cls['qualified_name'])
        
        self.stats['total_classes'] = len(self.class_index)
        self.stats['abstract_classes'] = len(self.abstract_classes)
        self.stats['final_classes'] = len(self.final_classes)
    
    def _extract_inheritance_relations(self, ast_data: Dict[str, Any]):
        """Extract inheritance relationships from AST data."""
        for relation_data in ast_data.get('inheritance_relations', []):
            relation = InheritanceRelation(
                derived_class=relation_data['derived_class'],
                base_class=relation_data['base_class'],
                access_level=relation_data.get('access_level', 'private'),
                is_virtual=relation_data.get('is_virtual', False),
                location=relation_data.get('location', {})
            )
            
            self.inheritance_relations.append(relation)
            
            # Update class info
            if relation.derived_class in self.class_hierarchy:
                self.class_hierarchy[relation.derived_class].base_classes.append(relation)
            
            if relation.base_class in self.class_hierarchy:
                self.class_hierarchy[relation.base_class].derived_classes.append(relation.derived_class)
        
        self.stats['inheritance_relations'] = len(self.inheritance_relations)
    
    def _build_inheritance_graphs(self):
        """Build forward and reverse inheritance graphs."""
        for relation in self.inheritance_relations:
            self.inheritance_graph[relation.derived_class].add(relation.base_class)
            self.reverse_inheritance_graph[relation.base_class].add(relation.derived_class)
    
    def _analyze_virtual_methods(self, ast_data: Dict[str, Any]):
        """Analyze virtual methods and their overrides."""
        # Index methods by class
        for func in ast_data.get('functions', []):
            if func.get('class_name'):
                self.method_index[func['class_name']].append(func)
        
        # Analyze virtual methods
        for class_name, methods in self.method_index.items():
            for method in methods:
                if method.get('is_virtual'):
                    virtual_method = VirtualMethodInfo(
                        name=method['name'],
                        qualified_name=method['qualified_name'],
                        declaring_class=class_name,
                        is_pure_virtual=method.get('is_pure_virtual', False),
                        is_final=method.get('is_final', False),
                        location=method['location']
                    )
                    
                    # Find overrides in derived classes
                    virtual_method.overriding_classes = self._find_method_overrides(
                        method['name'], class_name
                    )
                    
                    # Build override chain
                    virtual_method.override_chain = self._build_override_chain(
                        method['name'], class_name
                    )
                    
                    self.virtual_methods[method['qualified_name']] = virtual_method
                    
                    if virtual_method.is_pure_virtual:
                        self.stats['pure_virtual_methods'] += 1
                    else:
                        self.stats['virtual_methods'] += 1
    
    def _find_method_overrides(self, method_name: str, base_class: str) -> List[str]:
        """Find all classes that override a virtual method."""
        overriding_classes = []
        
        # Get all derived classes
        derived_classes = self._get_all_derived_classes(base_class)
        
        for derived_class in derived_classes:
            derived_methods = self.method_index.get(derived_class, [])
            for method in derived_methods:
                if (method['name'] == method_name and 
                    (method.get('is_virtual') or method.get('is_override'))):
                    overriding_classes.append(derived_class)
                    break
        
        return overriding_classes
    
    def _build_override_chain(self, method_name: str, class_name: str) -> List[str]:
        """Build the complete override chain for a virtual method."""
        chain = []
        
        # Start from the most derived class and work upward
        current_class = class_name
        visited = set()
        
        while current_class and current_class not in visited:
            visited.add(current_class)
            
            # Check if this class has the method
            class_methods = self.method_index.get(current_class, [])
            has_method = any(m['name'] == method_name for m in class_methods)
            
            if has_method:
                chain.append(current_class)
            
            # Move to base classes
            base_classes = list(self.inheritance_graph.get(current_class, []))
            if base_classes:
                current_class = base_classes[0]  # Take first base class
            else:
                break
        
        return list(reversed(chain))  # Return from base to derived
    
    def _get_all_derived_classes(self, base_class: str) -> Set[str]:
        """Get all classes derived from a base class (recursively)."""
        derived = set()
        queue = deque([base_class])
        visited = set()
        
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            
            immediate_derived = self.reverse_inheritance_graph.get(current, set())
            for derived_class in immediate_derived:
                if derived_class not in derived:
                    derived.add(derived_class)
                    queue.append(derived_class)
        
        return derived
    
    def _get_all_base_classes(self, derived_class: str) -> Set[str]:
        """Get all base classes of a derived class (recursively)."""
        bases = set()
        queue = deque([derived_class])
        visited = set()
        
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            
            immediate_bases = self.inheritance_graph.get(current, set())
            for base_class in immediate_bases:
                if base_class not in bases:
                    bases.add(base_class)
                    queue.append(base_class)
        
        return bases
    
    def _detect_diamond_inheritance(self):
        """Detect diamond inheritance patterns."""
        # Look for classes that have multiple paths to the same base class
        for class_name in self.class_index:
            base_classes = self._get_all_base_classes(class_name)
            
            # Check for duplicate base classes in different inheritance paths
            for base_class in base_classes:
                paths = self._find_inheritance_paths(class_name, base_class)
                if len(paths) > 1:
                    # Check if any path doesn't use virtual inheritance
                    non_virtual_paths = []
                    for path in paths:
                        if not self._is_virtual_inheritance_path(path):
                            non_virtual_paths.append(path)
                    
                    if non_virtual_paths:
                        diamond_case = {
                            'derived_class': class_name,
                            'common_base': base_class,
                            'inheritance_paths': paths,
                            'non_virtual_paths': non_virtual_paths,
                            'severity': 'error' if len(non_virtual_paths) > 1 else 'warning'
                        }
                        self.diamond_inheritance_cases.append(diamond_case)
        
        self.stats['diamond_inheritance_cases'] = len(self.diamond_inheritance_cases)
    
    def _find_inheritance_paths(self, derived_class: str, target_base: str) -> List[List[str]]:
        """Find all inheritance paths from derived to base class."""
        paths = []
        
        def dfs_paths(current: str, target: str, path: List[str], visited: Set[str]):
            if current == target:
                paths.append(path + [current])
                return
            
            if current in visited:
                return
            visited.add(current)
            
            # Explore base classes
            base_classes = self.inheritance_graph.get(current, set())
            for base_class in base_classes:
                dfs_paths(base_class, target, path + [current], visited.copy())
        
        dfs_paths(derived_class, target_base, [], set())
        return paths
    
    def _is_virtual_inheritance_path(self, path: List[str]) -> bool:
        """Check if an inheritance path uses virtual inheritance."""
        for i in range(len(path) - 1):
            derived = path[i]
            base = path[i + 1]
            
            # Find the inheritance relation
            for relation in self.inheritance_relations:
                if relation.derived_class == derived and relation.base_class == base:
                    if not relation.is_virtual:
                        return False
                    break
        return True
    
    def _build_vtables(self):
        """Build virtual function tables for classes."""
        for class_name, class_info in self.class_hierarchy.items():
            vtable = self._compute_vtable(class_name)
            self.vtables[class_name] = vtable
            class_info.vtable_size = len(vtable)
    
    def _compute_vtable(self, class_name: str) -> List[Dict[str, Any]]:
        """Compute the virtual function table for a class."""
        vtable = []
        
        # Get inheritance hierarchy (from most base to most derived)
        hierarchy = self._get_inheritance_hierarchy(class_name)
        hierarchy.append(class_name)  # Include the class itself
        
        # Collect virtual methods from hierarchy
        virtual_method_map = {}  # method_name -> method_info
        
        for cls in hierarchy:
            class_methods = self.method_index.get(cls, [])
            for method in class_methods:
                if method.get('is_virtual'):
                    method_name = method['name']
                    # Later methods override earlier ones
                    virtual_method_map[method_name] = {
                        'name': method_name,
                        'qualified_name': method['qualified_name'],
                        'declaring_class': cls,
                        'implementing_class': cls,
                        'is_pure_virtual': method.get('is_pure_virtual', False),
                        'vtable_index': len(virtual_method_map)
                    }
        
        # Convert to list (vtable order)
        for method_info in virtual_method_map.values():
            vtable.append(method_info)
        
        return vtable
    
    def _get_inheritance_hierarchy(self, class_name: str) -> List[str]:
        """Get the inheritance hierarchy from base to derived."""
        hierarchy = []
        visited = set()
        
        def dfs_hierarchy(current: str, path: List[str]):
            if current in visited:
                return
            visited.add(current)
            
            base_classes = list(self.inheritance_graph.get(current, []))
            if not base_classes:
                # This is a base class, start the hierarchy here
                hierarchy.extend(reversed(path + [current]))
                return
            
            for base_class in base_classes:
                dfs_hierarchy(base_class, path + [current])
        
        dfs_hierarchy(class_name, [])
        
        # Remove duplicates while preserving order
        seen = set()
        unique_hierarchy = []
        for cls in hierarchy:
            if cls not in seen:
                seen.add(cls)
                unique_hierarchy.append(cls)
        
        return unique_hierarchy[:-1] if unique_hierarchy else []  # Exclude the target class itself
    
    def _calculate_inheritance_depths(self):
        """Calculate inheritance depth for each class."""
        for class_name in self.class_hierarchy:
            depth = self._get_inheritance_depth(class_name)
            self.class_hierarchy[class_name].inheritance_depth = depth
        
        depths = [info.inheritance_depth for info in self.class_hierarchy.values()]
        if depths:
            self.stats['max_inheritance_depth'] = max(depths)
            self.stats['avg_inheritance_depth'] = sum(depths) / len(depths)
    
    def _get_inheritance_depth(self, class_name: str) -> int:
        """Get the inheritance depth of a class."""
        max_depth = 0
        visited = set()
        
        def dfs_depth(current: str, depth: int) -> int:
            if current in visited:
                return depth
            visited.add(current)
            
            base_classes = self.inheritance_graph.get(current, set())
            if not base_classes:
                return depth
            
            max_child_depth = depth
            for base_class in base_classes:
                child_depth = dfs_depth(base_class, depth + 1)
                max_child_depth = max(max_child_depth, child_depth)
            
            return max_child_depth
        
        return dfs_depth(class_name, 0)
    
    def _analyze_polymorphic_behavior(self, ast_data: Dict[str, Any]):
        """Analyze polymorphic method calls and behavior."""
        polymorphic_classes = set()
        
        # Analyze method calls for polymorphic behavior
        for call in ast_data.get('method_calls', []):
            object_type = call.get('object_type', {})
            method_name = call.get('method_name')
            
            if not method_name:
                continue
            
            # Get the static type of the object
            static_type = object_type.get('spelling', '').replace('*', '').replace('&', '').strip()
            
            if static_type in self.class_hierarchy:
                # Check if this method is virtual
                is_virtual_call = self._is_virtual_method_call(static_type, method_name)
                
                if is_virtual_call:
                    polymorphic_classes.add(static_type)
                    
                    # Find all possible runtime types
                    possible_types = self._get_possible_runtime_types(static_type)
                    
                    polymorphic_call = {
                        'location': call.get('location', {}),
                        'static_type': static_type,
                        'method_name': method_name,
                        'possible_runtime_types': list(possible_types),
                        'caller': call.get('caller_function', ''),
                        'dispatch_complexity': len(possible_types)
                    }
                    
                    if static_type in self.class_hierarchy:
                        self.class_hierarchy[static_type].polymorphic_calls.append(polymorphic_call)
        
        self.stats['polymorphic_classes'] = len(polymorphic_classes)
    
    def _is_virtual_method_call(self, class_name: str, method_name: str) -> bool:
        """Check if a method call is virtual."""
        class_methods = self.method_index.get(class_name, [])
        for method in class_methods:
            if method['name'] == method_name and method.get('is_virtual'):
                return True
        
        # Check base classes
        for base_class in self.inheritance_graph.get(class_name, []):
            if self._is_virtual_method_call(base_class, method_name):
                return True
        
        return False
    
    def _get_possible_runtime_types(self, static_type: str) -> Set[str]:
        """Get all possible runtime types for a static type."""
        possible_types = {static_type}
        
        # Add all derived classes (they could be the runtime type)
        derived_classes = self._get_all_derived_classes(static_type)
        possible_types.update(derived_classes)
        
        # Remove abstract classes (can't be instantiated)
        possible_types = {t for t in possible_types if t not in self.abstract_classes}
        
        return possible_types
    
    def _get_all_derived_classes(self, base_class: str) -> Set[str]:
        """Get all classes derived from a base class."""
        derived = set()
        queue = deque([base_class])
        visited = set()
        
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            
            immediate_derived = self.reverse_inheritance_graph.get(current, set())
            for derived_class in immediate_derived:
                if derived_class not in derived:
                    derived.add(derived_class)
                    queue.append(derived_class)
        
        return derived
    
    def get_class_hierarchy_info(self, class_name: str) -> Dict[str, Any]:
        """Get detailed hierarchy information for a specific class."""
        if class_name not in self.class_hierarchy:
            return {'error': f'Class {class_name} not found'}
        
        class_info = self.class_hierarchy[class_name]
        
        return {
            'class_name': class_name,
            'qualified_name': class_info.qualified_name,
            'is_abstract': class_info.is_abstract,
            'is_final': class_info.is_final,
            'inheritance_depth': class_info.inheritance_depth,
            'base_classes': [
                {
                    'name': rel.base_class,
                    'access_level': rel.access_level,
                    'is_virtual': rel.is_virtual
                }
                for rel in class_info.base_classes
            ],
            'derived_classes': class_info.derived_classes,
            'virtual_methods': class_info.virtual_methods,
            'pure_virtual_methods': class_info.pure_virtual_methods,
            'overridden_methods': class_info.overridden_methods,
            'polymorphic_calls': class_info.polymorphic_calls,
            'vtable_size': class_info.vtable_size,
            'vtable': self.vtables.get(class_name, []),
            'all_base_classes': list(self._get_all_base_classes(class_name)),
            'all_derived_classes': list(self._get_all_derived_classes(class_name))
        }
    
    def _get_all_base_classes(self, derived_class: str) -> Set[str]:
        """Get all base classes of a derived class."""
        bases = set()
        queue = deque([derived_class])
        visited = set()
        
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            
            immediate_bases = self.inheritance_graph.get(current, set())
            for base_class in immediate_bases:
                if base_class not in bases:
                    bases.add(base_class)
                    queue.append(base_class)
        
        return bases
    
    def resolve_virtual_call(self, object_type: str, method_name: str) -> List[Dict[str, Any]]:
        """
        Resolve a virtual method call to possible implementations.
        
        Args:
            object_type: Static type of the object
            method_name: Name of the method being called
            
        Returns:
            List of possible method implementations
        """
        possible_implementations = []
        
        # Get all possible runtime types
        runtime_types = self._get_possible_runtime_types(object_type)
        
        for runtime_type in runtime_types:
            class_methods = self.method_index.get(runtime_type, [])
            for method in class_methods:
                if method['name'] == method_name:
                    possible_implementations.append({
                        'class_name': runtime_type,
                        'method_name': method_name,
                        'qualified_name': method['qualified_name'],
                        'location': method['location'],
                        'is_override': method.get('is_override', False),
                        'is_final': method.get('is_final', False)
                    })
                    break
        
        return possible_implementations
    
    def _format_inheritance_results(self) -> Dict[str, Any]:
        """Format inheritance analysis results."""
        return {
            'class_hierarchy': {
                name: {
                    'name': info.name,
                    'qualified_name': info.qualified_name,
                    'location': info.location,
                    'is_abstract': info.is_abstract,
                    'is_final': info.is_final,
                    'base_classes': [
                        {
                            'name': rel.base_class,
                            'access_level': rel.access_level,
                            'is_virtual': rel.is_virtual,
                            'location': rel.location
                        }
                        for rel in info.base_classes
                    ],
                    'derived_classes': info.derived_classes,
                    'virtual_methods': info.virtual_methods,
                    'pure_virtual_methods': info.pure_virtual_methods,
                    'overridden_methods': info.overridden_methods,
                    'polymorphic_calls': info.polymorphic_calls,
                    'vtable_size': info.vtable_size,
                    'inheritance_depth': info.inheritance_depth
                }
                for name, info in self.class_hierarchy.items()
            },
            'inheritance_relations': [
                {
                    'derived_class': rel.derived_class,
                    'base_class': rel.base_class,
                    'access_level': rel.access_level,
                    'is_virtual': rel.is_virtual,
                    'location': rel.location
                }
                for rel in self.inheritance_relations
            ],
            'virtual_methods': {
                name: {
                    'name': vm.name,
                    'qualified_name': vm.qualified_name,
                    'declaring_class': vm.declaring_class,
                    'overriding_classes': vm.overriding_classes,
                    'is_pure_virtual': vm.is_pure_virtual,
                    'is_final': vm.is_final,
                    'override_chain': vm.override_chain,
                    'location': vm.location
                }
                for name, vm in self.virtual_methods.items()
            },
            'diamond_inheritance_cases': self.diamond_inheritance_cases,
            'vtables': self.vtables,
            'abstract_classes': list(self.abstract_classes),
            'final_classes': list(self.final_classes),
            'statistics': self.stats
        }